#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
script usage:

    python ssh.py -h
    python ssh.py -c cfgfile -l logfile

config json:
    [
      {
        "alias":    "aliasname",    # option, default user@host
        "user":     "root",         # option, default root
        "host":     "127.0.0.1",    # require
        "password": "iamgod",       # require if not identity
        "identity": "god.pem",      # require if not password
        "port":     22,             # option, default nil
        "shell":    "shell.sh",     # option, default nil
        "sudo":     false,          # option, default false
        "verbose":  false           # option, default false
      },
      ...
    ]
"""
import optparse
import os
import sys
import re
from os import path
from functools import partial

try:
    # https://github.com/pexpect/pexpect
    import pexpect
except ImportError, e:
    sys.exit("""Please install pexpect!

Install pexpect using pip:

    pip install pexpect

How to install pip:

    # https://pypi.python.org/pypi/setuptools/
    curl https://bootstrap.pypa.io/ez_setup.py -o - | sudo python
    sudo easy_install pip

    # https://pip.pypa.io/en/stable/installing/#install-pip
""")


class SSHError(Exception):
    def __init__(self, value='ssh error'):
        self.value = value
    def __str__(self):
        return str(self.value)
    def __repr__(self):
        return '<SSHError'

class SSHEOFError(SSHError):
    def __init__(self, value='ssh error: eof'):
        super(SSHEOFError, self).__init__(value)

class SSHTimeoutError(SSHError):
    def __init__(self, value='ssh error: timeout'):
        super(SSHTimeoutError, self).__init__(value)


class SSH(object):

    def __init__(self, **cfg):
        if not cfg.has_key('host'):
            raise SSHError('config error: host is required.\n'
                           'config value: ' + str(cfg))
        if not cfg.has_key('password') and not cfg.has_key('identity'):
            raise SSHError('config error: password or identity is required.\n'
                           'config value: ' + str(cfg))

        self.user = cfg.get('user', 'root')
        self.host = cfg.get('host')
        self.password = cfg.get('password', None)
        self.identity = cfg.get('identity', None)
        self.port = cfg.get('port', None)
        self.shell = cfg.get('shell', None)
        self.sudo = cfg.get('sudo', False)
        self.verbose = cfg.get('verbose', False)

        self.alias = cfg.get('alias', '%s@%s' % (self.user, self.host))

        self.logfile = cfg.get('logfile', sys.stdout)

        self.connected = False
        self.sudoed = False
        self._child = None

    def _spawn(self, timeout=5):
        cmd = '/usr/bin/ssh'
        args = []
        port = self.port
        if port != None:
            if not isinstance(port, basestring):
                port = str(port)
            args.append('-p')
            args.append(str(port))
        if not self.password and self.identity:
            args.append('-i')
            args.append(self.identity)
        args.append('%s@%s' % (self.user, self.host))
        # print("%s %s" % (cmd, ' '.join(args)))

        child = pexpect.spawn(cmd, args, timeout=timeout)
        # child.setecho(False)
        if self.verbose:
            child.logfile_read = self.logfile
            # child.logfile_send = None

        return child

    def _expect(self, child, pattern, timeout=-1):
        try:
            return child.expect(pattern, timeout=timeout)
        except pexpect.EOF:
            raise SSHEOFError
        except pexpect.TIMEOUT:
            raise SSHTimeoutError

    def _wait(self, child, timeout=0.5):
        child.expect([pexpect.EOF, pexpect.TIMEOUT], timeout=timeout)

    def connect(self):
        """
        # Scene: new key
        The authenticity of host '0.0.0.0 (0.0.0.0)' can't be established.
        RSA key fingerprint is SHA256:LnuJVH9/DXZDOiFtdJYHaojOIcvORuLnr0/IdbQdeuE.
        Are you sure you want to continue connecting (yes/no)? 

        # Scene: password
        root@0.0.0.0's password: 

        # Scene: success
        Welcome to ...
        Last login: ...

        # Scene: denied
        Permission denied, please try again.

        # Scene: refused
        ssh: connect to host 0.0.0.0 port 22: Connection refused

        # Scene: timeout
        ssh: connect to host 0.0.0.0 port 22: Operation timed out

        # Scene: no identity file
        Warning: Identity file *.pem not accessible: No such file or directory.
        Permission denied (publickey).
        """
        if self.connected:
            return

        HOSTKEY = 'Are you sure you want to continue connecting'
        PASSWORD = '[Pp]assword:'
        LOGIN = 'Last login:'
        DENIED = 'Permission denied'

        child = self._spawn()
        expect = partial(self._expect, child)
        sendline = child.sendline

        if self.password:
            i = expect([HOSTKEY, PASSWORD])
            if i == 0: # not public key, just accept it
                sendline('yes')
                expect([PASSWORD])
            sendline(self.password) # password
            i = expect([LOGIN, DENIED, PASSWORD])
            if i != 0:
                raise SSHError('ssh error: password incorrect')
        elif self.identity:
            i = expect([LOGIN, DENIED, HOSTKEY])
            if i == 2: # not public key, just accept it
                sendline('yes')
                i = expect([LOGIN, DENIED])
            if i != 0:
                raise SSHError('ssh error: identity incorrect')
        else:
            raise SSHError('ssh error: password or identity is required')

        self._child = child
        self.connected = True
        # wait read over, otherwise will output first send twice
        self._wait(child)

        if self.sudo:
            self.sendsudo()

    def send(self, shell=None):
        if not self.connected:
            raise SSHError('ssh error: not connected')
        if not shell:
            shell = self.shell
        if not shell:
            return
        sendline = self._child.sendline
        if path.exists(shell) and path.isfile(shell):
            with open(shell) as f:
                for line in f.xreadlines():
                    sendline(line.strip('\n'))
        elif isinstance(shell, basestring):
            sendline(shell)
        else:
            raise SSHError('ssh error: shell incorrect')
        self._wait(self._child)

    def sendsudo(self):
        if not self.connected:
            raise SSHError('ssh error: not connected')
        if self.sudoed:
            return

        PASSWORD = '[Pp]assword'
        ROOT = 'root'

        expect = partial(self._expect, self._child)
        sendline = self._child.sendline

        sendline('sudo -i')

        i = expect([ROOT, PASSWORD])
        if i == 1: # password
            if not self.password:
                raise SSHError('ssh error: sudo password is required')
            sendline(self.password)
            i = expect([ROOT, PASSWORD])
            if i != 0:
                raise SSHError('ssh error: sudo password incorrect')

        self.sudoed = True

    def sendfile(self):
        # scp [file] root@0.0.0.0:[dir]
        # scp -r [dir] root@0.0.0.0:[dir]
        pass

    def close(self):
        if not self._child:
            return
        child = self._child
        if child.isalive():
            if self.connected:
                child.sendline('exit')
                if self.sudoed:
                    self._wait(child)
                    child.sendline('exit')
                # wait read over and make output until eof
                self._wait(child)
            child.close()
        self._child = None
        self.sudoed = False
        self.connected = False


class SSHSaber(object):

    def __init__(self, logfile=sys.stdout, logprefix=None):
        self.logfile = logfile
        self.logprefix = logprefix

    def _log(self, s, prefix=None):
        # assert s != None
        if not isinstance(s, basestring): # not (str, unicode)
            s = str(s)
        log = self.logfile
        if prefix:
            s = s.replace('\n', '\n' + prefix)
            log.write(prefix)
        log.write(s)
        log.write(os.linesep)
        log.flush()

    def do(self, cfgdict):
        log = partial(self._log, prefix=self.logprefix)
        try:
            cfgdict['logfile'] = self.logfile
            ssh = SSH(**cfgdict)
        except SSHError as e:
            log(e)
            return
        try:
            log('ssh %s ...' % ssh.alias)
            ssh.connect()
            ssh.send()
            ssh.close()
            log('ssh %s success' % ssh.alias)
            ssh = None
        except SSHTimeoutError as e:
            if ssh.verbose:
                prefix = self.logprefix
                self._log(e, ('\n' + prefix) if prefix else '\n')
            else:
                log(e)
        except (SSHError, SSHEOFError) as e:
            log(e)
        finally:
            if ssh:
                ssh.close()
                log('ssh %s failed' % ssh.alias)

    def load(self, cfgfile):
        # from collections import namedtuple
        # hook = lambda d: namedtuple('X', d.keys())(*d.values())
        with open(cfgfile, 'r') as f:
            import json
            config = json.load(f)
            do_ssh = self.do
            if isinstance(config, list):
                for cfg in config:
                    do_ssh(cfg)
            elif isinstance(config, dict):
                do_ssh(config)
            else:
                raise ValueError('unknown json config format')


def _parse_args():
    parser = optparse.OptionParser()
    parser.add_option('-l', '--log', dest='logfile', default=None,
        help='log output file [default: %default]')
    parser.add_option('-c', '--cfg', dest='cfgfile', default='ssh.json',
        help='config json file [default: %default]')
    options, args = parser.parse_args()
    return options


if __name__ == "__main__":
    options = _parse_args()

    log = sys.stdout
    logprefix = '> '

    logfile = options.logfile
    if logfile:
        if not path.exists(logfile) or path.isfile(logfile):
            import codecs
            log = codecs.open(logfile, 'w', 'utf-8')
            # log = open(logfile, 'w')

    def close(o):
        o.flush()
        if hasattr(o, 'close'):
            o.close()

    # exit = sys.exit
    def exit(s):
        if logprefix:
            log.write(logprefix)
        log.write('error: ')
        log.write(s)
        log.write(os.linesep)
        close(log)
        sys.exit(1)

    cfgfile = options.cfgfile
    if not path.exists(cfgfile) or not path.isfile(cfgfile):
        exit('%s not an existed file' % cfgfile)

    saber = SSHSaber(log, logprefix)
    try:
        saber.load(cfgfile)
    except ValueError as e:
        exit('load error: %s' % e)

    close(log)
